package com.luo.ai.assistant.chat.impl;

import com.luo.ai.assistant.chat.ChatConsts;
import com.luo.ai.assistant.chat.ChatRagService;
import com.luo.ai.assistant.chat.MyChatRequest;
import com.luo.ai.assistant.chat.memory.MyDaoChatMemory;
import com.luo.ai.assistant.chat.memory.MyRedisChatMemory;
import com.luo.ai.assistant.chat.rag.MyQuestionAnswerAdvisor;
import com.luo.ai.assistant.chat.tool.MyChatTools;
import com.luo.ai.assistant.util.JsonUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.lang.Nullable;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.Objects;
import java.util.Optional;

/**
 * Chat with RAG (Retrieval Augmented Generation) 服务实现
 *
 * @author luohq
 * @date 2025-01-02 10:34
 */
@Service
public class ChatRagServiceImpl implements ChatRagService {

    private Logger log = LoggerFactory.getLogger(ChatRagServiceImpl.class);

    /**
     * 普通chat客户端
     */
    private final ChatClient chatClient;
    /**
     * RAG chat客户端
     */
    private final ChatClient chatRagClient;
    /**
     * chat记忆管理
     */
    private final ChatMemory chatMemory;
    private final MyDaoChatMemory myDaoChatMemory;

    public ChatRagServiceImpl(ChatClient.Builder chatClientBuilder,
                              VectorStore vectorStore,
                              MyDaoChatMemory myDaoChatMemory,
                              MyRedisChatMemory myRedisChatMemory,
                              //agent内置的工具
                              MyChatTools myChatTools,
                              //此处集成MCP Client获取的ToolCallbackProvider
                              //org.springframework.ai.mcp.AsyncMcpToolCallbackProvider
                              //https://github.com/spring-projects/spring-ai-examples/blob/main/model-context-protocol/client-starter/starter-webflux-client/src/main/java/org/springframework/ai/mcp/samples/client/Application.java#L38
                              @Nullable ToolCallbackProvider toolCallbackProvider) {
        //使用内存记忆聊天上下文（亦可使用CassandraChatMemory）
        //ChatMemory chatMemory = new InMemoryChatMemory();
        //使用DAO记忆聊天上下文
        this.chatMemory = myDaoChatMemory;
        this.myDaoChatMemory = myDaoChatMemory;
        //this.chatMemory = myRedisChatMemory;

        //普通chat客户端
        chatClientBuilder
                .defaultSystem(ChatConsts.DEFAULT_SYSTEM)
                .defaultAdvisors(
                        //记忆聊天上下文
                        MessageChatMemoryAdvisor.builder(this.chatMemory)
                                .conversationId(ChatConsts.DEFAULT_CHAT_MEMORY_CONVERSATION_ID)
                                .build(),
                        //logs the request and response data of the ChatClient
                        //注：logging.level.org.springframework.ai.chat.client.advisor=DEBUG
                        new SimpleLoggerAdvisor(ChatConsts.MY_REQUEST_TO_JSON_STRING, SimpleLoggerAdvisor.DEFAULT_RESPONSE_TO_STRING, 0)
                        //new SimpleLoggerAdvisor()
                );

        //若启用了MCP Client，则集成MCP Client获取的ToolCallbackProvider
        if (Objects.nonNull(toolCallbackProvider)) {
            //此处集成MCP Client获取的ToolCallbackProvider
            chatClientBuilder.defaultToolCallbacks(toolCallbackProvider);
        } else {
            //否则使用内置的工具
            chatClientBuilder.defaultTools(myChatTools);
        }

        this.chatClient = chatClientBuilder.build();

        //构建动态参数向量存储实现
        MyVectorStore myVectorStoreDynamic = new MyVectorStore(vectorStore);
        //RAG chat客户端
        this.chatRagClient = chatClientBuilder.clone()
                .defaultSystem(ChatConsts.DEFAULT_SYSTEM)
                .defaultAdvisors(
                        //记忆聊天上下文
                        MessageChatMemoryAdvisor.builder(this.chatMemory).conversationId(ChatConsts.DEFAULT_CHAT_MEMORY_CONVERSATION_ID).build(),
                        //RAG - 支持动态上设置相似度检索参数
                        MyQuestionAnswerAdvisor.builder(myVectorStoreDynamic)
                                .searchRequest(SearchRequest.builder().build())
                                .promptTemplate(PromptTemplate.builder().template(ChatConsts.DEFAULT_USER_TEXT_ADVISE).build())
                                .build(),
                        //logs the request and response data of the ChatClient
                        //注：logging.level.org.springframework.ai.chat.client.advisor=DEBUG
                        //默认order为0，这里设置为1，以便在RAG Advisor之后执行，以便记录RAG Advisor的输出
                        new SimpleLoggerAdvisor(ChatConsts.MY_REQUEST_TO_JSON_STRING, SimpleLoggerAdvisor.DEFAULT_RESPONSE_TO_STRING, 0 + 1)
                        //new SimpleLoggerAdvisor()
                )
                .build();
    }

    @Override
    public String chat(String userInput) {
        log.info("Chat input: {}", userInput);
        return this.chatClient.prompt()
                .user(userInput)
                //设置系统提示
                //.system("系统提示")
                //设置Tool工具链：https://docs.spring.io/spring-ai/reference/api/tools.html#_adding_tools_to_chatclient
                //.tools(myToolObj)
                //设置ToolContext：https://docs.spring.io/spring-ai/reference/api/tools.html#_tool_context
                //.toolContext(Map.of())
                .call()
                .content();
    }

    @Override
    public Flux<String> chatStream(String userInput) {
        log.info("Chat Stream input: {}", userInput);
        return this.chatClient.prompt()
                .user(userInput)
                .stream()
                .content();
    }

    @Override
    public String chatRag(String userInput) {
        log.info("Chat Rag input: {}", userInput);
        return this.chatRagClient.prompt()
                .user(userInput)
                .call()
                .content();
    }

    @Override
    public Flux<String> chatRagStream(String userInput) {
        log.info("Chat Rag Stream input: {}", userInput);
        return this.chatRagClient.prompt()
                .user(userInput)
                .stream()
                .content();
    }

    @Override
    public String chatRag(MyChatRequest myChatRequest) {
        log.info("Chat Rag Request: {}", JsonUtils.toJson(myChatRequest));
        return this.chatRagClient.prompt()
                .user(myChatRequest.getUserInput())
                .system(Optional.ofNullable(myChatRequest.getSystemInput()).orElse(ChatConsts.DEFAULT_SYSTEM))
                .advisors(a -> this.customAdvisor(a, myChatRequest))
                //运行时设置工具上下文参数
                .toolContext(Map.of(MyChatTools.USER_ID_PARAM, "luohq"))
                .call()
                .content();
    }

    @Override
    public Flux<String> chatRagStream(MyChatRequest myChatRequest) {
        log.info("Chat Rag Stream Request: {}", JsonUtils.toJson(myChatRequest));
        return this.chatRagClient.prompt()
                .user(myChatRequest.getUserInput())
                .system(Optional.ofNullable(myChatRequest.getSystemInput()).orElse(ChatConsts.DEFAULT_SYSTEM))
                .advisors(a -> this.customAdvisor(a, myChatRequest))
                //运行时设置工具上下文参数
                .toolContext(Map.of(MyChatTools.USER_ID_PARAM, "luohq"))
                .stream()
                .content();
    }

    /**
     * 自定义Advisor
     *
     * @param a             AdvisorSpec
     * @param myChatRequest MyChatRequest
     */
    private void customAdvisor(ChatClient.AdvisorSpec a, MyChatRequest myChatRequest) {
        //设置chat ID
        if (StringUtils.hasText(myChatRequest.getChatConvId())) {
            log.info("Chat RAG - Set Conversation ID: {}", myChatRequest.getChatConvId());
            a.param(ChatMemory.CONVERSATION_ID, myChatRequest.getChatConvId());
        }
        //设置chat历史记录数
        if (Objects.nonNull(myChatRequest.getChatHistoryCount())) {
            log.info("Chat RAG - Set History Count: {}", myChatRequest.getChatHistoryCount());
            //a.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY, myChatRequest.getChatHistoryCount());
            this.myDaoChatMemory.setLastN(myChatRequest.getChatHistoryCount());
        }



        //设置RAG过滤表达式
        if (StringUtils.hasText(myChatRequest.getRagFilterExpression())) {
            log.info("Chat RAG - Set Filter Expression: {}", myChatRequest.getRagFilterExpression());
            a.param(QuestionAnswerAdvisor.FILTER_EXPRESSION, myChatRequest.getRagFilterExpression());
        }
        //设置RAG相似度阈值
        if (Objects.nonNull(myChatRequest.getRagSimilarityThreshold())) {
            log.info("Chat RAG - Set Similarity Threshold: {}", myChatRequest.getRagSimilarityThreshold());
            a.param(MyQuestionAnswerAdvisor.SIMILARITY_THRESHOLD, myChatRequest.getRagSimilarityThreshold());
        }
        //设置RAG TopK
        if (Objects.nonNull(myChatRequest.getRagTopK())) {
            log.info("Chat RAG - Set TopK: {}", myChatRequest.getRagTopK());
            a.param(MyQuestionAnswerAdvisor.TOP_K, myChatRequest.getRagTopK());
        }
    }

}
